import torch
import torch.nn as nn
import pickle # Keep pickle if getTranslationModel is kept, otherwise remove

# Imports from other files in the same standalone_eval directory
import ModelInterfaces as mi # Changed from 'from ModelInterfaces import IASRModel'
import AIModels # Changed from 'from AIModels import NeuralASR'
import whisper_wrapper # Changed from 'from whisper_wrapper import WhisperASRModel'

def getASRModel(language: str, use_whisper:bool=True, device: str = None) -> mi.IASRModel:
    actual_device = None
    if device:
        actual_device = torch.device(device)
    else:
        actual_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if use_whisper:
        return whisper_wrapper.WhisperASRModel(device=actual_device) # Use actual_device
    
    # Silero STT part (alternative ASR)
    # Ensure silero_device is also correctly set if this path is taken.
    # The original code in your attached file already passes `device` to torch.hub.load
    # which expects a torch.device object or string. So, actual_device is fine.

    silero_lang_map = {
        'de': 'de',
        'en': 'en',
        'fr': 'fr' # Example, add other silero supported languages if needed
    }
    if language in silero_lang_map:
        model, decoder, utils = torch.hub.load(repo_or_dir='snakers4/silero-models',
                                               model='silero_stt',
                                               language=silero_lang_map[language],
                                               device=actual_device) # Use actual_device
        model.eval()
        return AIModels.NeuralASR(model, decoder) # Changed from NeuralASR
    else:
        raise ValueError(f'Language {language} not implemented for non-Whisper ASR or not supported by Silero.')

# getTTSModel and getTranslationModel are likely not needed for the core pronunciation evaluation demo.
# They can be removed if the standalone demo doesn't use TTS or translation.
# For now, I will keep them but comment out their content if they have external deps not being copied.

def getTTSModel(language: str) -> nn.Module:
    # This function depends on torch.hub.load for silero_tts
    # Ensure network access or pre-downloaded models if kept.
    print("TTS Model loading (getTTSModel) is part of the standalone files but might require network.")
    actual_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    silero_speaker_map = {
        'de': 'thorsten_v2',
        'en': 'lj_16khz'
    }
    if language in silero_speaker_map:
        speaker = silero_speaker_map[language]
        model, _ = torch.hub.load(repo_or_dir='snakers4/silero-models',
                                  model='silero_tts',
                                  language=language,
                                  speaker=speaker,
                                  device=actual_device) # Added device here too
        return model
    else:
        raise ValueError(f'TTS for language {language} not implemented')


def getTranslationModel(language: str):
    # This function depends on transformers and fetching models from Hugging Face Hub.
    # This will require network access.
    # Also, it saves pickled models to the current directory, which might be an issue for a clean standalone package.
    print("Translation Model loading (getTranslationModel) requires network and has side effects (pickling).")
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
    
    if language == 'de':
        model_name = "Helsinki-NLP/opus-mt-de-en"
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # Pickling might not be ideal for a standalone script, consider removing or making optional.
        # with open('translation_model_de.pickle', 'wb') as handle:
        #     pickle.dump(model, handle)
        # with open('translation_tokenizer_de.pickle', 'wb') as handle:
        #     pickle.dump(tokenizer, handle)
        return model, tokenizer
    else:
        raise ValueError(f'Translation for language {language} not implemented') 